import h5py
import torch
from torch.utils.data import DataLoader, Dataset


def transform_data(data, label, reverse=False):
    if reverse:
        label = 2 ** (label + 8)
        # ref_imgs [-1, 1] to [0, c]
        data = (data + 1) / 2 * label
    else:
        # normalize data [0, label] -> [-1, 1] # TODO check if this is correct
        data = (data / label) * 2 - 1
        # transform labels to integers, log2 - 8
        label = (torch.log2(label) - 8).int()
    return data, label


class H5Dataset(Dataset):
    def __init__(self, h5_file_path, preload=True):
        # Open the HDF5 file
        self.h5_file = h5py.File(h5_file_path, "r")

        # Assuming data is stored in 'data' and labels in 'labels'
        self.data = self.h5_file["showers"]  # shape [121000, 368]
        self.labels = self.h5_file["incident_energies"]  # shape [121000, 1]

        if preload:
            # Preload data into memory
            self.data = self.data[:]  # type: ignore
            self.labels = self.labels[:]  # type: ignore
            self.h5_file.close()

    def __len__(self):
        # Return the size of the dataset
        return len(self.data)  # type: ignore

    def __getitem__(self, idx):
        # Get the data and label at the specified index
        data = self.data[idx]  # type: ignore
        label = self.labels[idx]  # type: ignore

        # Convert to torch tensors
        label = torch.tensor(label, dtype=torch.int32)
        data = torch.tensor(data, dtype=torch.float32)
        return transform_data(data, label)


def get_data_loader(type, batch_size, shuffle, num_workers):
    if type == "1-phontons":
        dataset = H5Dataset("../data/dataset_1_photons_1.hdf5", preload=True)
        return DataLoader(
            dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers
        )
    else:
        raise ValueError("Invalid dataset type")


if __name__ == "__main__":
    loader = get_data_loader("1-phontons", batch_size=64, shuffle=True, num_workers=4)
    for batch in loader:
        x, y = batch
        print(x.shape, y.shape)
        print(x.dtype, y.dtype)
        break
    print(torch.unique(y))
    print(x.min(), x.max())
